import os
from typing import Annotated, List, Literal

from langchain_community.chat_models.zhipuai import ChatZhipuAI
from langchain_core.messages import BaseMessage
from langchain_core.tools import tool
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import add_messages, StateGraph, START, END
from langgraph.prebuilt import create_react_agent, ToolNode
from pydantic import BaseModel
from typing_extensions import TypedDict

os.environ["ZHIPUAI_API_KEY"] = "97738d4998b8732d707daf91a2b1c56d.2y6VKEuOlidwHDpI"

"""定义State开始"""


class State(TypedDict):
    messages: Annotated[List[BaseMessage], add_messages]
    summary: str
    intention: str


"""定义State结束"""

"""定义Tools开始"""


# 查询保单
@tool
def insurance_order_search(name: str, id_number: str):
    """根据用户的姓名和身份证号查询用户的保单"""
    return "{'123': '保险1', '456': '保险2', '789': '保险3'}"


# 保单退保
@tool
def insurance_order_cancel(name: str, id_number: str, policy_no: str):
    """根据用户提供的保单号进行退保操作"""
    pass


class AskHuman(BaseModel):
    """Ask the human a question"""
    question: str


tools = [insurance_order_search, insurance_order_cancel, AskHuman]
tool_node = ToolNode(tools)
"""定义Tools结束"""

"""定义Model开始"""
model = ChatZhipuAI(
    model='glm-4',
    temperature=0.5
)
agent_executor = create_react_agent(model, tool_node, state_modifier="""你是一个保险客户专员，现在已经进入到了客户退保的流程，请你按照如下步骤引导用户操作。
1、在退保的最开始阶段，要生成委婉的话术询问用户退保原因以及是否确认退保，并且说明保险对用户的好处。
2、如果用户确认要退保的话，需要优先询问用户，收集用户的姓名和身份证号后4位。如果用户确认不退保，则询问用户是否还有其他问题，如果没有其他问题的话，生成结束语给用户，并完结对话。
3、根据用户提供的姓名和身份证号，查询用户目前在途的保单信息，并提供给用户进行确认退保的保单号。
4、用户确认完退保的保单号之后，执行用户退保的操作。
5、执行完用户的退保操作之后，询问用户是否还有其他问题，如果没有其他问题的话，生成结束语给用户，并完结对话。
""", interrupt_before=["tools"])
"""定义Model结束"""

"""定义Nodes开始"""


class intention(TypedDict):
    intention: Literal["退保", "非退保其他保险业务", "结束会话"]


intention_recognition_prompt = """你是一个保险客户专员，能够结合用户的历史消息，识别出最新一条用户消息的核心意图。
历史消息总结:
{summary}

消息上下文
{context}

用户最新的消息
{user_message}
"""


def intention_recognition(state: State):
    prompt = format_prompt(state, intention_recognition_prompt)
    response = model.with_structured_output(intention).invoke(prompt)
    return {"intention": "其他" if response is None else response["intention"]}


def intention_condition(state: State):
    if state["intention"] == '退保':
        return "surrender"
    elif state["intention"] == '非退保其他保险业务':
        return "non_surrender"
    elif state["intention"] == "其他":
        return "salutation"
    else:
        return "end_session"


def format_prompt(state: State, prompt_template: str) -> str:
    context_messages = []
    if len(state["messages"]) > 1:
        context_messages = state["messages"][0:len(state["messages"]) - 1]
    context = ""
    for item in context_messages:
        context = context + item.type + ":" + item.content + "\n"
    prompt = prompt_template.format(
        summary=state.get("summary", "暂无"),
        context=context,
        user_message=state["messages"][len(state["messages"]) - 1].type + ":" + state["messages"][
            len(state["messages"]) - 1].content)
    return prompt


def non_surrender(state: State):
    return {"messages": [("ai", "好的，我帮你转人工客服进行处理")]}


def salutation(state: State):
    return {"messages": [("ai", "你好，我是水滴保险客服小依，有什么可以帮助到您的")]}


def end_session(state: State):
    return {"messages": [("ai", "好的，那我就先完结此次会话了，祝您生活愉快")]}


"""定义Nodes结束"""

"""定义WorkFlow开始"""
work_flow = StateGraph(State)
work_flow.add_node("intention_recognition", intention_recognition)
work_flow.add_node("agent_executor", agent_executor)
work_flow.add_node("non_surrender", non_surrender)
work_flow.add_node("end_session", end_session)
work_flow.add_node("salutation", salutation)

work_flow.add_edge(START, "intention_recognition")
work_flow.add_conditional_edges("intention_recognition",
                                intention_condition,
                                {"surrender": "agent_executor",
                                 "non_surrender": "non_surrender",
                                 "salutation": "salutation",
                                 "end_session": "end_session"})
work_flow.add_edge("non_surrender", END)
work_flow.add_edge("end_session", END)
work_flow.add_edge("salutation", END)
work_flow.add_edge("agent_executor", END)

graph = work_flow.compile(checkpointer=MemorySaver())

"""定义WorkFlow结束"""
